"""
FastMCP i18n - 优化版

运行方式:
    python i18n_mcp.py
"""

from mcp.server.fastmcp import FastMCP
import os
import json
from datetime import datetime
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass, asdict
from pathlib import Path
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class TranslationEntry:
    """国际化词条结构体"""
    chinese: str  # 中文内容
    i18n_key: str  # 国际化key
    translation: str = ""  # 翻译内容（可空，待翻译时为空）

class I18NStorage:
    """状态存储管理类（优化版）"""
    def __init__(self, state_file: str = "i18n_state.json"):
        self.state_file = state_file
        self._state = self._load_or_init_state()
    
    def _load_or_init_state(self) -> Dict[str, Any]:
        """加载或初始化状态文件"""
        if os.path.exists(self.state_file):
            try:
                with open(self.state_file, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                logger.error(f"加载状态文件失败: {e}")
        
        # 初始化默认状态
        return {
            "pending_files": [],
            "completed_files": [],
            "current_file": None,
            "entries": [],
            "updated_at": datetime.now().isoformat()
        }
    
    def _save_state(self) -> None:
        """保存状态到文件"""
        self._state["updated_at"] = datetime.now().isoformat()
        try:
            with open(self.state_file, 'w', encoding='utf-8') as f:
                json.dump(self._state, f, indent=2, ensure_ascii=False)
        except Exception as e:
            logger.error(f"保存状态文件失败: {e}")
    
    def add_files_to_queue(self, files: List[str]) -> int:
        """添加文件到待处理队列，返回新增文件数量"""
        existing_files = set(self._state["pending_files"] + self._state["completed_files"])
        new_files = [f.lower() for f in files if f not in existing_files]
        self._state["pending_files"].extend(new_files)
        self._save_state()
        return len(new_files)
    
    def get_next_file_to_process(self) -> Optional[str]:
        """获取下一个待处理文件"""
        if not self._state["pending_files"]:
            return None
        
        next_file = self._state["pending_files"][0]
        self._state["current_file"] = next_file
        self._save_state()
        return next_file
    
    def mark_file_completed(self, file_path: str) -> None:
        """将文件标记为已完成"""
        if file_path in self._state["pending_files"]:
            self._state["pending_files"].remove(file_path)
        if file_path not in self._state["completed_files"]:
            self._state["completed_files"].append(file_path)
        
        if self._state["current_file"] == file_path:
            self._state["current_file"] = None
        
        self._save_state()
    
    def create_entries(self, entries: List[TranslationEntry]) -> int:
        """添加词条（自动去重），返回新增词条数量"""
        existing_keys = {e["i18n_key"] for e in self._state["entries"]}
        new_entries = [
            asdict(entry) for entry in entries 
            if entry.i18n_key not in existing_keys
        ]
        
        if new_entries:
            self._state["entries"].extend(new_entries)
            self._save_state()
        
        return len(new_entries)
    
    def get_stats(self) -> Dict[str, int]:
        """获取统计信息"""
        return {
            "pending": len(self._state["pending_files"]),
            "completed": len(self._state["completed_files"]),
            "entries": len(self._state["entries"])
        }

class FileProcessor:
    """文件处理工具类"""
    
    # 支持的文件扩展名
    SUPPORTED_EXT = ('.java', '.js', '.ts', '.vue', '.jsx', '.tsx', '.py', '.txt')
    
    # 需要跳过的目录
    SKIP_DIRS = ('node_modules', '.git', 'dist', 'build', '__pycache__', '.venv')
    
    @classmethod
    def get_all_files(cls, dir_path: str) -> List[str]:
        """扫描目录下的所有目标文件"""
        if not os.path.isdir(dir_path):
            return []
        
        dir_path = os.path.abspath(dir_path)
        file_list = []
        
        for root, dirs, files in os.walk(dir_path):
            # 跳过指定目录
            dirs[:] = [d for d in dirs if d not in cls.SKIP_DIRS]
            
            # 收集支持的文件类型
            for file in files:
                if file.endswith(cls.SUPPORTED_EXT):
                    full_path = os.path.join(root, file)
                    file_list.append(full_path)
        
        return file_list
    
    @staticmethod
    def read_file_content(file_path: str, max_size: int = 50000) -> Tuple[str, bool]:
        """读取文件内容，返回内容和是否截断的标记"""
        try:
            file_size = os.path.getsize(file_path)
            if file_size > max_size:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read(max_size)
                return content, True
            
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read(), False
                
        except Exception as e:
            logger.error(f"读取文件失败 {file_path}: {e}")
            raise
    
    @staticmethod
    def write_file_content(file_path: str, content: str) -> bool:
        """将内容写入文件"""
        try:
            # 确保目录存在
            os.makedirs(os.path.dirname(file_path), exist_ok=True)
            
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(content)
            return True
        except Exception as e:
            logger.error(f"写入文件失败 {file_path}: {e}")
            return False

# 初始化MCP和存储
mcp = FastMCP("i18n-mcp")
storage = I18NStorage()

# 工具1：规划待处理文件
@mcp.tool()
async def plan(dir_path: str) -> str:
    """
    扫描目录，整理需要做国际化的相关文件路径，形成待处理列表
    
    参数:
        dir_path: 处理代码文件的目录（绝对路径）
    
    返回: JSON字符串，包含状态和统计信息
    """
    if not os.path.isdir(dir_path):
        return json.dumps({
            "status": "error",
            "message": f"目录不存在：{dir_path}"
        }, ensure_ascii=False)
    
    try:
        all_files = FileProcessor.get_all_files(dir_path)
        new_files_count = storage.add_files_to_queue(all_files)
        stats = storage.get_stats()
        
        return json.dumps({
            "status": "success",
            "total_files": len(all_files),
            "new_files_added": new_files_count,
            "pending_files": stats["pending"],
            "message": f"发现 {len(all_files)} 个文件，新增 {new_files_count} 个待处理文件"
        }, ensure_ascii=False, indent=2)
    
    except Exception as e:
        logger.error(f"规划文件失败: {e}")
        return json.dumps({
            "status": "error",
            "message": f"处理目录时发生错误: {str(e)}"
        }, ensure_ascii=False)

# 工具2：获取待处理文件
@mcp.tool()
async def todo_edit() -> str:
    """
    从待处理列表中取出一个文件进行国际化处理
    
    返回: JSON字符串，包含文件信息和处理提示
    """
    try:
        next_file = storage.get_next_file_to_process()
        if not next_file:
            return json.dumps({
                "isComplete": True,
                "message": "所有文件已处理完毕"
            }, ensure_ascii=False)
        
        # 读取文件内容
        content, truncated = FileProcessor.read_file_content(next_file)
        
        prompt = """
        请对这个文件进行国际化处理，将中文内容替换为国际化变量。
        你必须完成以下操作：
        1. 找出文件中所有需要展示给用户的中文文本（非注释、非代码逻辑中的中文）
        2. 将这些中文全部替换为国际化key（格式示例：{i18n.模块.功能}）
        3. 提取所有被替换的中文和对应key的映射关系
        
        ## 强制要求：
        1. 必须修改对应文件中的用户可见中文文本，替换为国际化key
        2. 必须返回完整的修改后代码（包含所有未修改部分，格式与原文件完全一致）
        3. 同时提取被替换的中文与对应key（用于后续存储）
        """
        
        instruction = """
        请严格按照以下格式返回结果（必须包含修改后的完整代码）：
        {
            "modified_code": "完整的修改后代码",
            "entries": [
                {"chinese": "中文内容1", "i18n_key": "module.key1", "translation": ""},
                {"chinese": "中文内容2", "i18n_key": "module.key2", "translation": ""}
            ]
        }
        """
        
        return json.dumps({
            "isComplete": False,
            "file": next_file,
            "content": content,
            "content_truncated": truncated,
            "prompt": prompt,
            "instruction": instruction
        }, ensure_ascii=False, indent=2)
    
    except Exception as e:
        logger.error(f"获取待处理文件失败: {e}")
        return json.dumps({
            "error": True,
            "message": f"处理文件时发生错误: {str(e)}"
        }, ensure_ascii=False)

# 工具3：保存处理结果
@mcp.tool()
async def save_results(modified_code: str, entries: List[dict], file_path: str) -> str:
    """
    保存国际化处理结果，包括修改后的代码和提取的词条
    
    参数:
        modified_code: 修改后的完整代码
        entries: 词条列表
        file_path: 文件路径
    
    返回: JSON字符串，包含处理结果和统计信息
    """
    try:
        # 验证参数
        if not modified_code or not isinstance(entries, list) or not file_path:
            return json.dumps({
                "status": "error",
                "message": "参数不完整或格式错误"
            }, ensure_ascii=False)
        
        # 保存修改后的代码到原文件
        success = FileProcessor.write_file_content(file_path, modified_code)
        if not success:
            return json.dumps({
                "status": "error",
                "message": f"保存修改后的代码到文件失败: {file_path}"
            }, ensure_ascii=False)
        
        # 转换并保存词条
        translation_entries = []
        for item in entries:
            if not isinstance(item, dict):
                continue
            
            # 验证必填字段
            if "chinese" not in item or "i18n_key" not in item:
                continue
                
            entry = TranslationEntry(
                chinese=item["chinese"],
                i18n_key=item["i18n_key"],
                translation=item.get("translation", "")
            )
            translation_entries.append(entry)
        
        # 保存词条
        saved_count = storage.create_entries(translation_entries)
        
        # 标记文件为已完成
        storage.mark_file_completed(file_path)
        
        # 获取统计信息
        stats = storage.get_stats()
        
        return json.dumps({
            "status": "success",
            "saved_entries": saved_count,
            "total_entries": stats["entries"],
            "remaining_files": stats["pending"],
            "message": f"成功处理文件 {os.path.basename(file_path)}，保存了 {saved_count} 个新词条"
        }, ensure_ascii=False, indent=2)
    
    except Exception as e:
        logger.error(f"保存处理结果失败: {e}")
        return json.dumps({
            "status": "error",
            "message": f"保存结果时发生错误: {str(e)}"
        }, ensure_ascii=False)

if __name__ == "__main__":
    mcp.run(transport="stdio")